# -*- coding: utf-8 -*-
# auther gym
from os import system, environ
environ["PBR_VERSION"] = "4.0.2"

from os.path import dirname, basename, join, isfile, getctime
import torch
import numpy as np
from utils.micaps import Micaps
from utils.shift import Shift
from model.reviseNet import Net
from utils.config import *
from tendo.singleton import SingleInstance
from glob import glob
import time


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# device = torch.device('cpu')

re_logger = Logger(filename=f'{logpath}/revise_infer_{datetime.now().year}-{datetime.now().month}.log', level='info',
                   when='D', back_count=30).logger


def revise(file):
    Tcls = basename(dirname(file))
    re_logger.info(f"Found scmoc file: {file}")
    mi = Micaps()
    _x = mi.read(file)
    if _x is None:
        re_logger.error("DATA ERROR! MV To issue path!")
        system(f"mv {file} {scmoc_issue_path}/{Tcls}")
        return
    _x[_x > Tem_max] = Tem_max
    _x[_x < Tem_min] = Tem_min
    # _x = _x[::-1, :]

    x = _x[:width, :height] + 273.15
    x_batch = torch.from_numpy(x[np.newaxis, np.newaxis, ...].copy())
    with torch.no_grad():
        out = model(x_batch.to(device=device, dtype=torch.float32))
        r_batch = out.cpu().numpy().squeeze()
    # r_batch = model.Predict(x_batch)
    r = r_batch.squeeze()
    _r = _x.copy()
    _r[:width, :height] = r - 273.15
    # _r = _r[::-1, :]

    _r[_r > Tem_max] = Tem_max
    _r[_r < Tem_min] = Tem_min

    revise_outdir = f"{revise_out_path}/{Tcls}"

    mi.write(f"{revise_outdir}/{basename(file)}", _r)

    if Tcls == 'TMAX':
        sh.WriteMax(mi, basename(file), _r)
    elif Tcls == 'TMIN':
        sh.WriteMin(mi, basename(file), _r)

    backPath = join(scmoc_backup_path, Tcls)
    command = 'mv %s %s' % (join(scmoc_new_path, Tcls, file), backPath)
    re_logger.info(f"{Tcls} {basename(file)} revise successed!")
    system(command)


if __name__ == "__main__":
    run = True

    try:
        me = SingleInstance()
        re_logger.info("revise_infer start...")

    except:
        re_logger.warning('revise infer process already  exist! ')
        run = False

    starttime = datetime.now()
    sh = Shift(f'{revise_out12_path}')

    while run:
        for tcls in Tcategorys:

            for ht in htimes:
                sfiles = sorted(glob(f'{scmoc_new_path}/{tcls}/*{ht}.???'))
                if len(sfiles) == 0:
                    continue

                weight_file = f"{weight_root}/{tcls}/{last_train_date}_{tcls}_{ht}.pth"
                re_logger.info(f"LOADING WEIGHTS: {weight_file} ")
                if not isfile(weight_file):
                    re_logger.error(f"{weight_file} Not Found")

                model = Net()
                checkpoints = torch.load(weight_file, map_location=device)
                state_dict = checkpoints['state_dict']
                model.load_state_dict(state_dict)
                model.to(device)
                model.eval()

                for sfile in sfiles:
                    if time.time() - getctime(sfile) < 0.1:
                        continue
                    revise(sfile)

                torch.cuda.empty_cache()
                starttime = datetime.now()

        if (datetime.now() - starttime).seconds > process_duration:
            re_logger.info(f"Timeout, Idle {process_duration}, revise_infer exits! ")
            break
